package com.king;

import javassist.ClassPool;
import javassist.CtClass;
import javassist.CtMethod;

import java.io.IOException;
import java.lang.instrument.Instrumentation;
import java.nio.file.Files;
import java.nio.file.Paths;
import java.util.ArrayList;
import java.util.List;

public class NacosServiceDiscoveryAgent {

    private static List<String> configLines = new ArrayList<>();

    public static void premain(String agentArgs, Instrumentation inst) {
        System.out.println("加载自定义NacosServiceDiscoveryAgent处理");
        // 读取配置文件
        if (agentArgs != null && !agentArgs.isEmpty()) {
            loadConfigFromFile(agentArgs);
        } else {
            // 尝试从当前jar包目录读取配置文件
            String jarFilePath = NamingGrpcClientProxyAgent.class.getProtectionDomain().getCodeSource().getLocation().getPath();
            String jarDir = jarFilePath.substring(0, jarFilePath.lastIndexOf('/') + 1);
            String configPath = jarDir + "nacosConfig.txt";
            if (configPath.startsWith("/")){
                configPath = configPath.replaceFirst("/", "");
            }
            loadConfigFromFile(configPath);
        }

        if (configLines == null || configLines.isEmpty()) {
            System.out.println("NacosAgent未发现配置文件nacosConfig.txt,取消Agent处理");
            return;
        }

        inst.addTransformer((loader, className, classBeingRedefined, protectionDomain, classfileBuffer) -> {
            // 解析配置
            StringBuilder configMapCode = new StringBuilder();
                    configMapCode.append("   java.util.Map configMap = new java.util.HashMap(); \n");
            for (int i = 1; i < configLines.size(); i++) {
                String[] splitStr = configLines.get(i).split(" ");
                if (splitStr.length == 3){
                    configMapCode.append("   configMap.put(\"").append(splitStr[0]).append("\", \"").append(splitStr[1]).append(",").append(splitStr[2]).append("\"); \n");
                }
            }

            // 第一处
            String classNameTarget = "com.alibaba.cloud.nacos.discovery.NacosServiceDiscovery";
            if (classNameTarget.replace(".", "/").equals(className)) {
                // 使用Javassist库对类进行转换
                javassist.ClassPool pool = javassist.ClassPool.getDefault();
                try {
                    javassist.CtClass cc = pool.get(classNameTarget);
                    javassist.CtMethod m = cc.getDeclaredMethod("getInstances");

                    // 动态代码块
                    StringBuilder codeBlock = new StringBuilder();
                    // 加载自定义配置
                    codeBlock.append(configMapCode);
                    codeBlock.append(
//                            "          java.util.List instances = getInstances($$); \n" +
                            "          java.util.List instances = $_; \n" +
                            "          if (instances != null && !instances.isEmpty()) {\n" +
                            "            int i = 0; \n" +
                            "            while (i < instances.size()) { \n" +
                            "                org.springframework.cloud.client.ServiceInstance instance = instances.get(i);\n" +
                            "                if (instance instanceof com.alibaba.cloud.nacos.NacosServiceInstance) {\n" +
                            "                   com.alibaba.cloud.nacos.NacosServiceInstance ins = (com.alibaba.cloud.nacos.NacosServiceInstance) instance;\n" +
                            "                   if (configMap.containsKey(ins.getServiceId())) {\n" +
                            "                       String configStr = configMap.get(ins.getServiceId()); \n" +
                            "                       String[] split = configStr.split(\",\"); \n" +
                            "                       System.out.println(\"服务:\" + ins.getServiceId() + \" 替换实例IP端口：\" + ins.getHost() + \":\" + instance.getPort() + \" -> \" + configStr);\n" +
                            "                       ins.setHost(split[0]);\n" +
                            "                       ins.setPort(Integer.parseInt(split[1])); \n" +
                            "                   }\n" +
                            "                }\n" +
                            "                i = i + 1; \n" +
                            "            } \n" +
                            "          }\n" +
                            "       return instances;");
                    String block = codeBlock.toString();
                    System.out.println("代理代码块NacosServiceDiscovery.getInstances：\n" + block);
//                    m.insertBefore(block);
                    m.insertAfter(block);
                    byte[] bytecode = cc.toBytecode();
                    cc.detach();
                    return bytecode;
                } catch (Exception e) {
                    System.out.println("加载异常");
                    e.printStackTrace();
                }
            }

            // 打印调用节点信息
            try {
                // 只拦截 Feign 的  类
                if ("org/springframework/cloud/loadbalancer/blocking/client/BlockingLoadBalancerClient".equals(className)) {
                    return transformLoadBalancerFeignClient(classBeingRedefined);
                }
            } catch (Exception e) {
                e.printStackTrace();
            }

            return classfileBuffer;
        });
    }

    private static byte[] transformLoadBalancerFeignClient(Class<?> classBeingRedefined) throws Exception {
        ClassPool pool = ClassPool.getDefault();
        CtClass ctClass = pool.get("org.springframework.cloud.loadbalancer.blocking.client.BlockingLoadBalancerClient");

        // 拦截 execute 方法
        CtClass[] params = new CtClass[2];
        params[0] = pool.get("java.lang.String");
        params[1] = pool.get("org.springframework.cloud.client.loadbalancer.Request");

        CtMethod executeMethod = ctClass.getDeclaredMethod("choose", params);
        executeMethod.insertAfter("org.springframework.cloud.client.ServiceInstance instance = $_; " +
                "if (instance != null) {" +
                "   System.out.println(\"服务:\" + instance.getServiceId() + \" 调用实例IP端口：\" + instance.getHost() + \":\" + instance.getPort());\n" +
                "}" +
                "return instance;");
        byte[] bytecode = ctClass.toBytecode();
        ctClass.detach();
        return bytecode;
    }

    private static void loadConfigFromFile(String filePath) {
        try {
            List<String> lines = Files.readAllLines(Paths.get(filePath));
            configLines.addAll(lines);
        } catch (IOException e) {
            System.err.println("Failed to read config file: " + filePath);
            e.printStackTrace();
        }
    }

    public static List<String> getConfigLines() {
        return configLines;
    }
}
